////////////////////////////////////////////////////////////////////////////////
// Includes
////////////////////////////////////////////////////////////////////////////////

// Header
#include "ANNPlugin.hpp"

// Random number generator
#include <cstdlib>

// Math
#include <cmath>

// Types
#include "VFloat.hpp"

// Conversion functions
#include "Convert.hpp"

// I/O
#include <iostream>

// Wildcard definition
#include "SCProgram.hpp"

// ANN Macros
#include "ANNMacros.hpp"

// Toolkit macros
#include "TKMacros.h"

// Namespaces
using namespace std;
using namespace Types;
using namespace Toolkit;
using namespace SC;

////////////////////////////////////////////////////////////////////////////////
// Global Network Variables
////////////////////////////////////////////////////////////////////////////////

// Learning rate
#define L_RATE				0.5 //005

// Momentum rate
#define M_RATE				0.0

// Noise rate
#define N_RATE				0.0
#define N_TH				1.0

// Maximum weight value
#define MAX_WEIGHT			7.5

// Maximum initial weight value
#define MAX_INIT_WEIGHT		0.5

// Sigmoide function slope value
#define SIG_SLOPE			2.0

////////////////////////////////////////////////////////////////////////////////
// Tool Functions
////////////////////////////////////////////////////////////////////////////////

// Zero in the unsigned integer string format
static const string zeroS(INTEGER_SIZE,'0');

// Zero in the following VFloat format
static const VFloat zeroF(IS_SIGNED, INT_P_SIZE, DEC_P_SIZE, 0.);

// Maximum float value
static const double max_float = (pow(2.,(double)INT_P_SIZE) - 1.) + ((pow(2.,(double)DEC_P_SIZE)-1.)/pow(2.,(double)DEC_P_SIZE));

// No (or Error) value
static const VFloat noValF(IS_SIGNED, INT_P_SIZE, DEC_P_SIZE, max_float);

// Generic names of the transfer function and its derivative
#define transfer sigmoid
#define transfer_prime sigmoid_prime

// Transfer function
inline const double sigmoid(const double x, const double sl = SIG_SLOPE,  const double sh = 0.) {
	return (2. / (1 + exp(- sl * (x - sh)))) - 1.;
}
/*
inline const double sigmoid01(const double x, const double sl = SIG_SLOPE,  const double sh = 0.) {
	return 1. / (1 + exp(- sl * (x - sh)));
}
*/

// First derivative of the transfer function
inline const double sigmoid_prime(const double x, const double sl = SIG_SLOPE,  const double sh = 0.) {
	return (2.*sl * exp(- sl * (x - sh))) / pow(1 + exp(- sl * (x - sh)), 2);
}
/*
inline const double sigmoid01_prime(const double x, const double sl = SIG_SLOPE,  const double sh = 0.) {
	return (sl * exp(- sl * (x - sh))) / pow(1 + exp(- sl * (x - sh)), 2);
}
*/

////////////////////////////////////////////////////////////////////////////////
// Plugin init and exit
////////////////////////////////////////////////////////////////////////////////

void InitPlugin (const string & prog_name) {
#ifdef PLUGIN_DEBUG
	cout.precision(12);
	cerr.precision(12);
	cout << "VFloat bounds are [ " << -max_float << " ; " << max_float << " ]" << endl;
#endif
}

void ClosePlugin () {
}

////////////////////////////////////////////////////////////////////////////////
// SC Callbacks
////////////////////////////////////////////////////////////////////////////////

bool InputTransfer (SCSystem & universe, const SCSystem & input_transfer, SCSystem & axon) {
	bool has_init = false;
	// Get the universe input data size 
	int nb_inputs = BStrToUInt(universe.GetSchema1(UN_S1_INPUT_SIZE_IDX, INTEGER_SIZE));
	// If no data, no computation can happen
	if (nb_inputs == 0) return false;
	// Get the universe data to transfer
	int data_ref = BStrToUInt(axon.GetSchema2(IA_S2_DATA_REF_IDX, INTEGER_SIZE));
	/*
	// If no data has been allocated
	if (data_ref == 0) {
		// Get the universe number of already allocated data
		int nb_allocated = BStrToUInt(universe.GetSchema1(UN_S1_NB_ALLOCATED_IDX, INTEGER_SIZE));
		// Allocate if possible
		if (nb_allocated < nb_inputs) {
			// Allocate
			data_ref = nb_allocated + 1;
			//input_transfer.ReplaceKernel(UIntToBStr(data_ref, INTEGER_SIZE), IT_K_DATA_REF_IDX, INTEGER_SIZE);
			// Increment allocated data
			++nb_allocated;
			universe.ReplaceSchema1(UIntToBStr(nb_allocated, INTEGER_SIZE), UN_S1_NB_ALLOCATED_IDX, INTEGER_SIZE);
			// Some computation happened
			has_init = true;
		}
#ifdef PLUGIN_DEBUG
		else
			cout << "!!! No more input data to allocate, all " << nb_inputs << " are already allocated !!!" << endl ;
#endif
	}
	*/
	// Assert some data can be transferred (slot ok and new data)
	if (data_ref != 0) {
		// Get signal value
		string signal = universe.GetSchema1(UN_S1_INPUT_IDX + (data_ref-1)*FLOAT_SIZE, FLOAT_SIZE);
		if (signal != noValF.GetStringValue()) {
			// Set signal value
			axon.ReplaceSchema1(signal, AX_S1_SIGNAL_IDX, FLOAT_SIZE);
			// Set sum value (same here than the input)
			axon.ReplaceSchema1(signal, AX_S1_SUM_IDX, FLOAT_SIZE);
#ifdef PLUGIN_DEBUG
			VFloat vFloat(IS_SIGNED, INT_P_SIZE, DEC_P_SIZE);
			vFloat.StringValue() = signal;
			cout << "-> Set input : " << vFloat.GetFloatValue() << endl;
#endif
			// Reset universe new data
			universe.ReplaceSchema1(noValF.GetStringValue(), UN_S1_INPUT_IDX + (data_ref-1)*FLOAT_SIZE, FLOAT_SIZE);
			// If udpdate is required, increment data flag in axon (to signal that weights can be updated)
			if (universe.GetKernel(UN_K_UPDATE_FLAG_IDX, 1) == C_TRUE) {
#ifdef PLUGIN_DEBUG
				cout << "# Set flag for update" << endl;
#endif
				axon.ReplaceKernel(
					UIntToBStr( (BStrToUInt(axon.GetKernel(AX_K_DFLAG_IDX, DFLAG_LGH)) + 1) % (1<<DFLAG_LGH), DFLAG_LGH),
					AX_K_DFLAG_IDX, DFLAG_LGH
				);
				universe.ReplaceKernel(C_FALSE, UN_K_UPDATE_FLAG_IDX, 1);
			}
			// Computation happened
			return true;
		}
	}
	// No computation happened but eventually the initialisation
	return has_init;
}

bool Synapse (SCSystem & axon, SCSystem & synapse, SCSystem & dendrites) {
	VFloat vFloat(IS_SIGNED, INT_P_SIZE, DEC_P_SIZE);
	// --- Get synapse slot ---
	int slot = BStrToUInt(synapse.GetKernel().substr(SY_K_SLOT_IDX, INTEGER_SIZE));
	// If no slot exists yet, create it
	if (slot <= 0) {
#ifdef PLUGIN_DEBUG
		cout << "# Initialising synapse" << endl ;
#endif
		int nb_de_data = BStrToUInt(dendrites.GetSchema1(DE_S1_INPUT_SIZE_IDX, INTEGER_SIZE));
		// See whether an old slot can be reused
		slot = 1;
		VFloat vFloat(IS_SIGNED, INT_P_SIZE, DEC_P_SIZE);
		// If a slot is free, the loop will stop, otherwise slot will be a new data
		while (slot <= nb_de_data) {
			if (dendrites.GetSchema1(DE_S1_INPUT_IDX + ((slot-1)*FLOAT_SIZE), FLOAT_SIZE) == noValF.GetStringValue())
				break;
			++slot;
		}
		synapse.ReplaceKernel(UIntToBStr(slot, INTEGER_SIZE), SY_K_SLOT_IDX, INTEGER_SIZE);
#ifdef PLUGIN_DEBUG
		cout << "-> Creating synapse dendritic slot : " << slot << endl ;
#endif
		// Increase dendrites data size
		if (slot > nb_de_data) {
			++nb_de_data;
			dendrites.ReplaceSchema1(UIntToBStr(nb_de_data, INTEGER_SIZE), DE_S1_INPUT_SIZE_IDX, INTEGER_SIZE);
		}
		// Initialise signal to zero
		dendrites.ReplaceSchema1(zeroF.GetStringValue(), DE_S1_INPUT_IDX + (slot-1)*FLOAT_SIZE, FLOAT_SIZE);
		// Initialise weight
		vFloat.SetStringFromFloat( (2. * MAX_INIT_WEIGHT * frand()) - MAX_INIT_WEIGHT );
		synapse.ReplaceKernel(vFloat.GetStringValue(), SY_K_WEIGHT_IDX, FLOAT_SIZE);
#ifdef PLUGIN_DEBUG
		cout << "-> Initialise synapse weight " << vFloat.GetFloatValue() << endl;
#endif
		// Initialise error
		dendrites.ReplaceSchema2(zeroF.GetStringValue(), DE_S2_ERRORS_IDX + (slot-1)*FLOAT_SIZE, FLOAT_SIZE);
	}
#ifdef PLUGIN_DEBUG
	cout << "-> Found slot : " << slot << endl;
#endif

	// Get axon signal
	vFloat.StringValue() = axon.GetSchema1(AX_S1_SIGNAL_IDX, FLOAT_SIZE);
	double signal = vFloat.GetFloatValue();

	// --- Error Backpropagation and Delta Weight Update ---
	// Get error of the ahead neuron (information corresponding to this synapse)
	vFloat.StringValue() = dendrites.GetSchema2(DE_S2_ERRORS_IDX + (slot-1)*FLOAT_SIZE, FLOAT_SIZE);
	double err_den = vFloat.GetFloatValue();
	// Set this ahead error to zero
	dendrites.ReplaceSchema2(zeroF.GetStringValue(), DE_S2_ERRORS_IDX + (slot-1)*FLOAT_SIZE, FLOAT_SIZE);
	// Get weight
	vFloat.StringValue() = synapse.GetKernel().substr(SY_K_WEIGHT_IDX, FLOAT_SIZE);
	double weight = vFloat.GetFloatValue();
	// Get sum value in axon
	vFloat.StringValue() = axon.GetSchema1(AX_S1_SUM_IDX, FLOAT_SIZE);
	double sum = vFloat.GetFloatValue();
	// Compute synapse error
	double err_syn = err_den * weight * transfer_prime(sum);
	// Get axon error
	vFloat.StringValue() = axon.GetSchema1(AX_S1_ERR_IDX, FLOAT_SIZE);
	double err_ax = vFloat.GetFloatValue();
	// Add synapse error to axon error limiting error to max_float
	err_ax += err_syn;
	BOUND(err_ax,-max_float,max_float);
	vFloat.SetStringFromFloat(err_ax);
	axon.ReplaceSchema1(vFloat.GetStringValue(), AX_S1_ERR_IDX, FLOAT_SIZE);
	// Update delta weight
	vFloat.StringValue() = synapse.GetKernel().substr(SY_K_D_WEIGHT_IDX, FLOAT_SIZE);
	double d_weight = vFloat.GetFloatValue();
	d_weight += L_RATE * err_den * signal  + ((ABS(err_den)>=N_TH)?(N_RATE*(2.*frand()-1.)):0.f);
	BOUND(d_weight,-max_float,max_float);
	vFloat.SetStringFromFloat(d_weight);
	synapse.ReplaceKernel(vFloat.GetStringValue(), SY_K_D_WEIGHT_IDX, FLOAT_SIZE);

	// --- Potential Weight Update ---
	if (BStrToUInt(axon.GetKernel(AX_K_DFLAG_IDX, DFLAG_LGH)) == ((BStrToUInt(synapse.GetKernel(SY_K_DFLAG_IDX, DFLAG_LGH))+1)%(1<<DFLAG_LGH)) ) {
		synapse.ReplaceKernel(axon.GetKernel(AX_K_DFLAG_IDX, DFLAG_LGH), SY_K_DFLAG_IDX, DFLAG_LGH);
		dendrites.ReplaceKernel(axon.GetKernel(AX_K_DFLAG_IDX, DFLAG_LGH), DE_K_DFLAG_IDX, DFLAG_LGH);
		// Old delta weight
		vFloat.StringValue() = synapse.GetKernel().substr(SY_K_OLD_D_WEIGHT_IDX, FLOAT_SIZE);
		double old_d_weight = vFloat.GetFloatValue();
#ifdef PLUGIN_DEBUG
		cout << "-> Update weight " << weight;
#endif
		weight += d_weight + M_RATE*old_d_weight;
		BOUND(weight,-MAX_WEIGHT,MAX_WEIGHT);
		vFloat.SetStringFromFloat(weight);
#ifdef PLUGIN_DEBUG
		cout << " to " << vFloat.GetFloatValue() << endl;
#endif
		synapse.ReplaceKernel(vFloat.GetStringValue(), SY_K_WEIGHT_IDX, FLOAT_SIZE);
		old_d_weight = d_weight + M_RATE*old_d_weight;
		BOUND(old_d_weight,-max_float,max_float);
		vFloat.SetStringFromFloat(old_d_weight);
		synapse.ReplaceKernel(vFloat.GetStringValue(), SY_K_OLD_D_WEIGHT_IDX, FLOAT_SIZE);
		synapse.ReplaceKernel(zeroF.GetStringValue(), SY_K_D_WEIGHT_IDX, FLOAT_SIZE);
	}

	// --- Signal Transfer ---
	// Store result = signal * weight
	signal = weight * signal;
	BOUND(signal,-max_float,max_float);
	vFloat.SetStringFromFloat(signal);
	dendrites.ReplaceSchema1(vFloat.GetStringValue(), DE_S1_INPUT_IDX + ((slot-1)*FLOAT_SIZE), FLOAT_SIZE);

	// Computation happenned
	return true;
}

bool Membrane (SCSystem & dendrites, SCSystem & membrane, SCSystem & axon) {
	VFloat vFloat(IS_SIGNED, INT_P_SIZE, DEC_P_SIZE);
	// Get the dendrites data size 
	int nb_de_data = BStrToUInt(dendrites.GetSchema1(DE_S1_INPUT_SIZE_IDX, INTEGER_SIZE));

	// Initialise if needed the threshold
	if (membrane.GetKernel().substr(ME_K_THRESHOLD_IDX, FLOAT_SIZE) == noValF.GetStringValue()) {
		vFloat.SetStringFromFloat( (2. * MAX_INIT_WEIGHT * frand()) - MAX_INIT_WEIGHT );
		membrane.ReplaceKernel(vFloat.GetStringValue(),ME_K_THRESHOLD_IDX, FLOAT_SIZE);
#ifdef PLUGIN_DEBUG
		cout << "-> Initialise membrane threshold " << vFloat.GetFloatValue() << endl;
#endif
	}

	// --- Error Transfer ---
	// Get the output error for learning
	vFloat.StringValue() = axon.GetSchema1(AX_S1_ERR_IDX, FLOAT_SIZE);
	double err_ax = vFloat.GetFloatValue();
#ifdef PLUGIN_DEBUG
	if (err_ax != 0.)
		cout << "-> Error transfer " << err_ax << endl;
#endif
	// Set axon error to zero
	axon.ReplaceSchema1(zeroF.GetStringValue(), AX_S1_ERR_IDX, FLOAT_SIZE);
	// Add error to dendrites (one value for each synapse)
	for (int d=0; d<nb_de_data; ++d) {
		vFloat.StringValue() = dendrites.GetSchema2(DE_S2_ERRORS_IDX + (d*FLOAT_SIZE), FLOAT_SIZE);
		double err_den = vFloat.GetFloatValue();
		err_den += err_ax;
		BOUND(err_den,-max_float,max_float);
		vFloat.SetStringFromFloat(err_den);
		dendrites.ReplaceSchema2(vFloat.GetStringValue(), DE_S2_ERRORS_IDX + (d*FLOAT_SIZE), FLOAT_SIZE);
	}

	// --- Delta Threshold Update ---
	vFloat.StringValue() = membrane.GetKernel().substr(ME_K_D_THRESHOLD_IDX, FLOAT_SIZE);
	double d_threshold = vFloat.GetFloatValue();
	d_threshold -= L_RATE * err_ax + ((ABS(err_ax)>=N_TH)?(N_RATE*(2.*frand()-1.)):0.f);
	BOUND(d_threshold,-max_float,max_float);
	vFloat.SetStringFromFloat(d_threshold);
	membrane.ReplaceKernel(vFloat.GetStringValue(), ME_K_D_THRESHOLD_IDX, FLOAT_SIZE);

	// Get threshold
	vFloat.StringValue() = membrane.GetKernel().substr(ME_K_THRESHOLD_IDX, FLOAT_SIZE);
	double threshold = vFloat.GetFloatValue();

	// --- Potential Threshold Update ---
	if (BStrToUInt(dendrites.GetKernel(DE_K_DFLAG_IDX, DFLAG_LGH)) == ((BStrToUInt(membrane.GetKernel(ME_K_DFLAG_IDX, DFLAG_LGH))+1)%(1<<DFLAG_LGH)) ) {
		membrane.ReplaceKernel(dendrites.GetKernel(DE_K_DFLAG_IDX, DFLAG_LGH), ME_K_DFLAG_IDX, DFLAG_LGH);
		axon.ReplaceKernel(dendrites.GetKernel(DE_K_DFLAG_IDX, DFLAG_LGH), AX_K_DFLAG_IDX, DFLAG_LGH);
		// Old delta threshold
		vFloat.StringValue() = membrane.GetKernel().substr(ME_K_OLD_D_THRESHOLD_IDX, FLOAT_SIZE);
		double old_d_threshold = vFloat.GetFloatValue();
#ifdef PLUGIN_DEBUG
		cout << "-> Update threshold " << threshold;
#endif
		threshold += d_threshold + M_RATE*old_d_threshold;
		BOUND(threshold,-MAX_WEIGHT,MAX_WEIGHT);
		vFloat.SetStringFromFloat(threshold);
#ifdef PLUGIN_DEBUG
		cout << " to " << vFloat.GetFloatValue() << endl;
#endif
		membrane.ReplaceKernel(vFloat.GetStringValue(), ME_K_THRESHOLD_IDX, FLOAT_SIZE);
		old_d_threshold = d_threshold + M_RATE*old_d_threshold;
		BOUND(old_d_threshold,-max_float,max_float);
		vFloat.SetStringFromFloat(old_d_threshold);
		membrane.ReplaceKernel(vFloat.GetStringValue(), ME_K_OLD_D_THRESHOLD_IDX, FLOAT_SIZE);
		membrane.ReplaceKernel(zeroF.GetStringValue(), ME_K_D_THRESHOLD_IDX, FLOAT_SIZE);
	}

	// --- Signal Transfer ---
	// Sum the data
#ifdef PLUGIN_DEBUG
	cout << "-> Membrane (size " << nb_de_data << ") : transform ( ";
#endif
	// Sum computation
	double sum = 0.f;
	for (int d=0; d<nb_de_data; ++d) {
		vFloat.StringValue() = dendrites.GetSchema1(DE_S1_INPUT_IDX + (d*FLOAT_SIZE), FLOAT_SIZE);
		sum += vFloat.GetFloatValue();
#ifdef PLUGIN_DEBUG
		cout << vFloat.GetFloatValue() << " ";
#endif
	}
#ifdef PLUGIN_DEBUG
	cout << ")" << endl;
#endif
	// Substract threshold from the sum
	sum -= threshold;
	BOUND(sum,-max_float,max_float);
	// Store sum in the axon
	vFloat.SetStringFromFloat(sum);
	axon.ReplaceSchema1(vFloat.GetStringValue(), AX_S1_SUM_IDX, FLOAT_SIZE);
	// Apply transfer function
	double signal = transfer(sum);
	// Store signal in the axon
	vFloat.SetStringFromFloat(signal);
	axon.ReplaceSchema1(vFloat.GetStringValue(), AX_S1_SIGNAL_IDX, FLOAT_SIZE);
#ifdef PLUGIN_DEBUG
	cout << " into " << vFloat.GetFloatValue() << endl ;
#endif	

	// Computation happenned
	return true;
}

bool OutputTransfer (SCSystem & axon, const SCSystem & output_transfer, SCSystem & universe) {
	VFloat vFloat(IS_SIGNED, INT_P_SIZE, DEC_P_SIZE);
	// Get the universe slot where to transfer
	int slot = BStrToUInt(axon.GetSchema2(OA_S2_DATA_REF_IDX, INTEGER_SIZE));
	// Ensure the universe number of outputs is consistent
	int nb_outputs = BStrToUInt(universe.GetSchema2(UN_S2_OUTPUT_SIZE_IDX, INTEGER_SIZE));
	if (slot > nb_outputs) {
#ifdef PLUGIN_DEBUG
		cout << "Setting ouput size to : " << slot << endl ;
#endif
		// Set output data size
		nb_outputs = slot;
		universe.ReplaceSchema2(UIntToBStr(nb_outputs, INTEGER_SIZE), UN_S2_OUTPUT_SIZE_IDX, INTEGER_SIZE);
	}

	//BStrToUInt(output_transfer.GetKernel().substr(OT_K_SLOT_IDX, INTEGER_SIZE));
/*
	// If no slot has been allocated
	if (slot == 0) {
		// Get the universe output data size 
		int nb_outputs = BStrToUInt(universe.GetSchema2(UN_S2_OUTPUT_SIZE_IDX, INTEGER_SIZE));
		// Allocate
		slot = nb_outputs + 1;
		output_transfer.ReplaceKernel(UIntToBStr(slot, INTEGER_SIZE), OT_K_SLOT_IDX, INTEGER_SIZE);
#ifdef PLUGIN_DEBUG
		cout << "Creating ouput slot : " << slot << endl ;
#endif
		// Increment universe output data size
		++nb_outputs;
		universe.ReplaceSchema2(UIntToBStr(nb_outputs, INTEGER_SIZE), UN_S2_OUTPUT_SIZE_IDX, INTEGER_SIZE);
	}
	if (slot <= 0) {
#ifdef PLUGIN_DEBUG
		cerr << "Fatal error : Slot should be positive but is " << slot << endl;
#endif
		return false;
	}
*/
	// --- Supervision ---
	if (universe.GetKernel(UN_K_SUPERVISION_FLAG_IDX,1) == C_TRUE) {
		// Get axon signal
		vFloat.StringValue() = axon.GetSchema1(AX_S1_SIGNAL_IDX, FLOAT_SIZE);
		double signal = vFloat.GetFloatValue();
		// Get expected signal
		vFloat.StringValue() = universe.GetKernel(UN_K_EXPECTED_IDX + (slot-1)*FLOAT_SIZE, FLOAT_SIZE);
		double expected = vFloat.GetFloatValue();
		// Add error in axon
		vFloat.StringValue() = axon.GetSchema1(AX_S1_ERR_IDX, FLOAT_SIZE);
		double err_ax = vFloat.GetFloatValue();
		err_ax += (expected-signal);//*transfer_prime(sum);
		BOUND(err_ax,-max_float,max_float);
		vFloat.SetStringFromFloat(err_ax);
		axon.ReplaceSchema1(vFloat.GetStringValue(), AX_S1_ERR_IDX, FLOAT_SIZE);
#ifdef PLUGIN_DEBUG
		cout << "# Setting error in axon to " << vFloat.GetFloatValue() << endl ;
#endif
		// Set universe supervision flag to false
		universe.ReplaceKernel(C_FALSE, UN_K_SUPERVISION_FLAG_IDX, 1);
	}
	else
		axon.ReplaceSchema1(zeroF.GetStringValue(), AX_S1_ERR_IDX, FLOAT_SIZE);

	// --- Signal Transfer ---
	universe.ReplaceSchema2(axon.GetSchema1(AX_S1_SIGNAL_IDX, FLOAT_SIZE), UN_S2_OUTPUT_IDX + (slot-1)*FLOAT_SIZE, FLOAT_SIZE);
#ifdef PLUGIN_DEBUG
	vFloat.StringValue() = axon.GetSchema1(AX_S1_SIGNAL_IDX, FLOAT_SIZE);
	cout << "-> Output signal: " << vFloat.GetFloatValue() << endl;
#endif

	// Computation happened
	return true;
}

////////////////////////////////////////////////////////////////////////////////
